package security.crypto.aes;


/*
 *****************************************************************************
 *       AES Encryption.													 *
 *       																	 *
 *       Advanced Encryption Standard implementation in JAVA.   			 *
 *       																	 *
 *       Solution based on Niyaz PK's implementation in C.      			 *
 *		 All credits reserved.												 * 
 *																			 *
 *****************************************************************************
 *   This is the part of the application to encrypt binary files using 		 *
 *   AES algorithm. 														 *
 *****************************************************************************
 * @filename: AES_encrypt.java             	                                 *
 * @language: Java                                                           *
 * @last update at: 2010-04-13                                               *
 *****************************************************************************
 * @students that colaborated on this file:                                  *
 *  55962 - Manuel Monterroso - m.monterroso@gmail.com			             *
 *  57441 - Ant�nio Lages - antoniomslages@ist.utl.pt                        *
 *****************************************************************************
 */
class AES_encrypt {
	
	private byte[] rest = new byte[0];
	private final byte[] _initialIV = {(byte)0x19 ,(byte) 0xa5 ,(byte) 0xd6 ,(byte) 0x9d ,(byte) 0xc9 ,(byte) 0xa8 ,(byte) 0xe9 ,(byte) 0xea ,(byte) 0xe6 ,(byte)0x74 ,(byte) 0x1e ,(byte) 0xa5 ,(byte)0x39 ,(byte) 0x35 ,(byte) 0x73 ,(byte) 0xb7};
	private byte[] IV;
	
	private int cmode;
	private AES _aes;
	
	void init(AES aes){
		_aes = aes;
		cmode = _aes.getCMode();
		IV = new byte[16];
		System.arraycopy(_initialIV, 0, IV, 0, _initialIV.length);
	}
	
	// This method performs the ciphering of the input data, using the parameters specified in the init method.
	// Note: If the last bytes of the input data do not form a full block, the bytes are not ciphered at that instance.
	byte[] update (byte[] plaintext) {
		
		int i, j, round = 0;
		// state - the array that holds the intermediate results during encryption
		byte[][] state = new byte[4][4];
		
		int quotient = plaintext.length / 16; 
		int remainder = plaintext.length % 16;
		
		if (rest.length > 0) {
			
			byte[] newPlaintext = new byte[plaintext.length + rest.length];
				
			for(i=0; i < rest.length; i++)
				newPlaintext[i] = rest[i];
			
			for(i=0; i < plaintext.length; i++)
				newPlaintext[rest.length + i] = plaintext[i];
			
			quotient = newPlaintext.length / 16;
			remainder = newPlaintext.length % 16;	
			
			plaintext = newPlaintext;
			
			rest = new byte[0];
		}
		
		if (remainder != 0) {
			
			rest = new byte[remainder];
			
			for(i=0; i < remainder; i++)
				rest[i] = plaintext[quotient*16 + i];
		}
		
		int pt_length = quotient*16;
		
		for (int globalOffset=0; globalOffset < pt_length;) {
			
			switch(cmode) {
			
			case 2 : { // CTR mode with a constant initialization vector as nounce
				
				//Copy the input IV (as nounce) to state array			
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = IV[i*4 + j];
					}
				}
				
				// Add the first round key to the state before starting the rounds
				state = addRoundKey(state, _aes.getRoundKey(), 0);
				
				// There will be Nr rounds.
				// The first Nr-1 rounds are identical.
				// These Nr-1 rounds are executed in the loop below.
				for(round=1; round < _aes.getNr(); round++) {
					
					state = addRoundKey(mixColumns(shiftRows(subBytes(state))), _aes.getRoundKey(), round);
				}
				
				// The last round is given below
				// The MixColumns function is not here in the last round
				state = addRoundKey(shiftRows(subBytes(state)), _aes.getRoundKey(), round);

				// The encryption process is over. Copy the state array to output array.
				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						plaintext[globalOffset + i*4 + j] ^= state[j][i];
					}
				}
				
				break; 
			}
			
			default : { // ECB or CBC modes
				
				if (cmode==1) // CBC mode only
					for (i=0; i < 16; i++)
						plaintext[globalOffset + i] ^= IV[i];
				
				//Copy the input plaintext to state array			
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = plaintext[globalOffset + i*4 + j];
					}
				}
				
				// Add the first round key to the state before starting the rounds
				state = addRoundKey(state, _aes.getRoundKey(), 0);
				
				// There will be Nr rounds.
				// The first Nr-1 rounds are identical.
				// These Nr-1 rounds are executed in the loop below.
				for(round=1; round < _aes.getNr(); round++) {
					
					state = addRoundKey(mixColumns(shiftRows(subBytes(state))), _aes.getRoundKey(), round);
				}
				
				// The last round is given below
				// The MixColumns function is not here in the last round
				state = addRoundKey(shiftRows(subBytes(state)), _aes.getRoundKey(), round);

				// The encryption process is over. Copy the state array to output array.
				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						
						plaintext[globalOffset + i*4 + j] = state[j][i];
					}
				}
				
				if (cmode==1) // CBC mode only
					for (i=0; i < 16; i++)
						IV[i] = plaintext[globalOffset + i];
				
			}
			}
			
			globalOffset+=16;
		}
		
		return plaintext;
	}

	// This method performs the ciphering of the input data, using the parameters specified in the init method.
	// Note: If the last bytes of the input data do not form a full block, the bytes are not ciphered at that instance.
	byte[] update (byte[] plaintext, int inputlen) {
		
		int i, j, round = 0;
		// state - the array that holds the intermediate results during encryption
		byte[][] state = new byte[4][4];
		
		int quotient = inputlen / 16; 
		int remainder = inputlen % 16;
		
		if (rest.length > 0) {
			
			byte[] newPlaintext = new byte[inputlen + rest.length];
				
			for(i=0; i < rest.length; i++)
				newPlaintext[i] = rest[i];
			
			for(i=0; i < inputlen; i++)
				newPlaintext[rest.length + i] = plaintext[i];
			
			quotient = newPlaintext.length / 16;
			remainder = newPlaintext.length % 16;	
			
			plaintext = newPlaintext;
			
			rest = new byte[0];
		}
		
		if (remainder != 0) {
			
			rest = new byte[remainder];
			
			for(i=0; i < remainder; i++)
				rest[i] = plaintext[quotient*16 + i];
		}
		
		int pt_length = quotient*16;
		
		for (int globalOffset=0; globalOffset<pt_length;) {
			
			switch(cmode) {
			
			case 2 : { // CTR mode with a constant initialization vector as nounce
				
				//Copy the input IV (as nounce) to state array			
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = IV[i*4 + j];
					}
				}
				
				// Add the first round key to the state before starting the rounds
				state = addRoundKey(state, _aes.getRoundKey(), 0);
				
				// There will be Nr rounds.
				// The first Nr-1 rounds are identical.
				// These Nr-1 rounds are executed in the loop below.
				for(round=1; round < _aes.getNr(); round++) {
					
					state = addRoundKey(mixColumns(shiftRows(subBytes(state))), _aes.getRoundKey(), round);
				}
				
				// The last round is given below
				// The MixColumns function is not here in the last round
				state = addRoundKey(shiftRows(subBytes(state)), _aes.getRoundKey(), round);

				// The encryption process is over. Copy the state array to output array.
				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						plaintext[globalOffset + i*4 + j] ^= state[j][i];
					}
				}
				
				break; 
			}
			
			default : { // ECB or CBC modes
				
				if (cmode==1) // CBC mode only
					for (i=0; i < 16; i++)
						plaintext[globalOffset + i] ^= IV[i];
				
				//Copy the input plaintext to state array			
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = plaintext[globalOffset + i*4 + j];
					}
				}
				
				// Add the first round key to the state before starting the rounds
				state = addRoundKey(state, _aes.getRoundKey(), 0);
				
				// There will be Nr rounds.
				// The first Nr-1 rounds are identical.
				// These Nr-1 rounds are executed in the loop below.
				for(round=1; round < _aes.getNr(); round++) {
					
					state = addRoundKey(mixColumns(shiftRows(subBytes(state))), _aes.getRoundKey(), round);
				}
				
				// The last round is given below
				// The MixColumns function is not here in the last round
				state = addRoundKey(shiftRows(subBytes(state)), _aes.getRoundKey(), round);

				// The encryption process is over. Copy the state array to output array.
				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						
						plaintext[globalOffset + i*4 + j] = state[j][i];
					}
				}
				
				if (cmode==1) // CBC mode only
					for (i=0; i < 16; i++)
						IV[i] = plaintext[globalOffset + i];
				
			}
			}
			
			globalOffset+=16;
		}
		
		return plaintext;
	}
	
	// This method performs the ciphering of the input data, using the parameters specified in the init
	// method. This method also performs the necessary procedures to conclude the ciphering of the
	// data stream, including the addition of necessary padding.
	byte[] doFinal (byte[] plaintext) {
		
		int i; 
		
		int quotient = plaintext.length / 16; 
		int remainder = plaintext.length % 16;
		
		if (rest.length > 0) {
			
			byte[] newPlaintext = new byte[plaintext.length + rest.length];
				
			for(i=0; i < rest.length; i++)
				newPlaintext[i] = rest[i];
			
			for(i=0; i < plaintext.length; i++)
				newPlaintext[rest.length + i] = plaintext[i];
			
			quotient = newPlaintext.length / 16;
			remainder = newPlaintext.length % 16;	
			
			plaintext = newPlaintext;
			
			rest = new byte[0];
		}
		
		if (remainder != 0) {
			
			byte[] temp = new byte[remainder];
			byte[] block = new byte[16];
			
			for(i=0; i < remainder; i++)
				temp[i] = plaintext[quotient*16 + i];
			
			block = padding(temp);
			
			byte[] pt_temp = new byte[(quotient+1)*16];
			
			for(i=0; i < quotient*16; i++)
				pt_temp[i] = plaintext[i];
			
			for(i=0; i < 16; i++)
				pt_temp[quotient*16 + i] = block[i];
			
			plaintext = pt_temp;
			
		} else {
			
			byte[] pt_temp = new byte[plaintext.length + 16];
			
			for(i=0; i < quotient*16; i++)
				pt_temp[i] = plaintext[i];
			
			for(i=0; i < 16; i++)
				pt_temp[quotient*16 + i] = 16;

			plaintext = pt_temp;
		}
		
		return update (plaintext);
	}
	
	// This method performs the ciphering of the input data, using the parameters specified in the init
	// method. This method also performs the necessary procedures to conclude the ciphering of the
	// data stream, including the addition of necessary padding.	
	byte[] doFinal (byte[] plaintext, int inputlen) {
		
		int i;
		
		int quotient = inputlen / 16; 
		int remainder = inputlen % 16;
		
		if (rest.length > 0) {
			
			byte[] newPlaintext = new byte[inputlen + rest.length];
				
			for(i=0; i < rest.length; i++)
				newPlaintext[i] = rest[i];
			
			for(i=0; i < inputlen; i++)
				newPlaintext[rest.length + i] = plaintext[i];
			
			quotient = newPlaintext.length / 16;
			remainder = newPlaintext.length % 16;	
			
			plaintext = newPlaintext;
			
			rest = new byte[0];
		}
		
		if (remainder != 0) {
			
			byte[] temp = new byte[remainder];
			byte[] block = new byte[16];
			
			for(i=0; i < remainder; i++)
				temp[i] = plaintext[quotient*16 + i];
			
			block = padding(temp);
			
			byte[] pt_temp = new byte[(quotient+1)*16];
			
			for(i=0; i < quotient*16; i++)
				pt_temp[i] = plaintext[i];
			
			for(i=0; i < 16; i++)
				pt_temp[quotient*16 + i] = block[i];
			
			plaintext = pt_temp;
			
		} else {
			
			byte[] pt_temp = new byte[inputlen + 16];
			
			for(i=0; i < inputlen; i++)
				pt_temp[i] = plaintext[i];
			
			for(i=0; i < 16; i++)
				pt_temp[quotient*16 + i] = 16;

			plaintext = pt_temp;

		}	
		
		return update (plaintext);
	}
	
	//*****************************************************************************************************************	
	// The padding() function completes the last block to be called within the cipher function doFinal()
	static byte[] padding (byte[] block) {
		
		int length = block.length;
		int temp=0;
		
		if(length > 16) {
			System.out.println("Padding block bigger than 128 bits (16 bytes)!");
			System.exit(-1);
		} else 
			temp = 16 - length;
		
		byte[] b_temp = new byte[16];
		
		for(int i=0; i<length; i++)
			b_temp[i] = block[i];
		
		for(int i=0; i<temp; i++)
			b_temp[16-temp+i] = (byte) temp;
			
		return b_temp;
	}

	// The SubBytes() function substitutes the values in the
	// state matrix with values in an S-box.
	byte[][] subBytes (byte[][] state) {
		
		int i,j;
		
		for(i=0;i<4;i++) {
			
			for(j=0;j<4;j++) {
				
				state[i][j] = (byte)_aes.getSBoxValue(state[i][j]);
			}
		}
		
		return state;
	}

	// The ShiftRows() function shifts the rows in the state to the left.
	// Each row is shifted with different offset.
	// Offset = Row number. So the first row is not shifted.
	byte[][] shiftRows (byte[][] state) {
		
		byte temp;

		temp = state[1][0];
		state[1][0] = state[1][1];
		state[1][1] = state[1][2];
		state[1][2] = state[1][3];
		state[1][3] = temp;

		temp = state[2][0];
		state[2][0] = state[2][2];
		state[2][2] = temp;
		temp = state[2][1];
		state[2][1] = state[2][3];
		state[2][3] = temp;

		temp=state[3][0];
		state[3][0]=state[3][3];
		state[3][3]=state[3][2];
		state[3][2]=state[3][1];
		state[3][1]=temp;
		
		return state;
	}
	
	// This function adds the round key to state.
	// The round key is added to the state by an XOR function.
	byte[][] addRoundKey (byte[][] state, byte[] roundKey, int round) {
		
		int i,j;
		
		for(i=0; i < 4; i++) {
			
			for(j=0; j < 4; j++) {
				
				state[j][i] ^= roundKey[round * _aes.getNb() * 4 + i * _aes.getNb() + j];
			}
		}
		
		return state;
	}
	
	// MixColumns function mixes the columns of the state matrix
	byte[][] mixColumns (byte[][] state) {
		
		int i;
		byte tmp,tm,t;
		
		for(i=0; i < 4; i++) {	
			
			t = state[0][i];
			tmp = (byte) (state[0][i] ^ state[1][i] ^ state[2][i] ^ state[3][i]) ;
			tm = (byte) (state[0][i] ^ state[1][i]) ; tm = xtime(tm); state[0][i] ^= tm ^ tmp ;
			tm = (byte) (state[1][i] ^ state[2][i]) ; tm = xtime(tm); state[1][i] ^= tm ^ tmp ;
			tm = (byte) (state[2][i] ^ state[3][i]) ; tm = xtime(tm); state[2][i] ^= tm ^ tmp ;
			tm = (byte) (state[3][i] ^ t) ; tm = xtime(tm); state[3][i] ^= tm ^ tmp ;
		}
		
		return state;
	}
	
	// xtime is a function that finds the product of {02} and the argument to xtime modulo {1b}  
	byte xtime (byte x) {
		
		return (byte) ((x<<1) ^ (((x>>7) & 1) * 0x1b));
	}
}
